package com.ycz.algorithm.utils.tree.avl;

import sun.reflect.generics.tree.Tree;

/**
 * @author yancz
 * @ClassName AVLTree
 * @description: AVL树常用方法
 * @date 2025-01-02 15:13:08
 * @version: 1.0
 */
public class AVLTree {

    private static TreeNode root; // 根结点

    // 获取节点高度，空树的高度为-1
    int height(TreeNode node) {
        return node == null ? -1 : node.height;
    }

    // 更新节点高度，节点高度 = 最高子树高度 + 1
    void updateHeight(TreeNode node) {
        if (node != null) {
            node.height = Math.max(height(node.left), height(node.right)) + 1;
        }
    }

    // 计算树节点的平衡因子，平衡因子 = 左子树高度 - 右子树高度
    // AVL树的每个节点平衡因子只可能是-1、0、1
    int balanceFactor(TreeNode node) {
        if (node == null) return 0;
        return height(node.left) - height(node.right);
    }

    /*
     * @author yancz
     * @description 右旋操作
     * @date 2025/1/2 15:44
     * @param node 失衡节点
     * @return TreeNode
     **/
    TreeNode rightRotate(TreeNode node) {
        TreeNode child = node.left; // 失衡节点的左节点作为child
        TreeNode grandChild = child.right; // 失衡节点的右节点作为grandChild
        child.right = node; // child为原点，将node进行右旋
        node.left = grandChild;
        updateHeight(node);
        updateHeight(child);
        return child;
    }

    // 左旋操作，与右旋是对称操作
    TreeNode leftRotate(TreeNode node) {
        TreeNode child = node.right;
        TreeNode grandChild = child.left;
        child.left = node;
        node.right = grandChild;
        updateHeight(node);
        updateHeight(child);
        return child;
    }

    // 通用旋转操作
    TreeNode commonRotate(TreeNode node) {
        int balanceFactor = balanceFactor(node); // 获取失衡节点的平衡因子
        if (balanceFactor > 1) { // 左偏树
            if (balanceFactor(node.left) >= 0) { // 直接右旋
                return rightRotate(node);
            } else { // 先左旋后右旋
                node.left = leftRotate(node.left);
                return rightRotate(node);
            }
        }
        if (balanceFactor < -1) { // 右偏树
            if (balanceFactor(node.right) <= 0) { // 直接左旋
                return leftRotate(node);
            } else { // 先右旋后左旋
                node.right = rightRotate(node.right);
                return leftRotate(node);
            }
        }
        return node;
    }

    // 插入节点
    void insert(int val) {
        root = insertHelper(root, val);
    }

    /*
     * @author yancz
     * @description 递归插入节点
     * avl树插入节点后，从该节点到根结点的路径上可能会出现一系列失衡节点
     * 需要从插入节点开始，自底向上进行旋转，使得所有失衡节点重新恢复平衡
     * @date 2025/1/3 15:54
     * @param node
     * @param val
     * @return TreeNode
     **/
    private TreeNode insertHelper(TreeNode node, int val) {
        if (node == null) {
            return new TreeNode(val);
        }
        // 查找插入节点的位置
        if (val < node.val) { // 如果插入节点值<当前节点值
            node.left = insertHelper(node.left, val); // 递归寻找插入节点位置
        } else if (val > node.val) { // 如果插入节点值>当前节点值
            node.right = insertHelper(node.right, val);
        } else { // 如果插入节点值=当前节点值，不操作
            return node;
        }
        updateHeight(node); // 插入节点后需要更新节点高度
        node = commonRotate(node); // 调整失衡树，直到整棵树处于平衡
        return node;
    }

    // 删除节点
    void remove(int val) {
        root = removeHelper(root, val);
    }

    /*
     * @author yancz
     * @description // 递归删除节点
     * @date 2025/1/7 17:08
     * @param node
     * @param val
     * @return TreeNode
     */
    TreeNode removeHelper(TreeNode node, int val) {
        if (node == null) {
            return null;
        }
        if (val < node.val) { // 待删除节点值<当前节点值，左递归查找
            node.left = removeHelper(node.left, val);
        } else if (val > node.val) { // 待删除节点值>当前节点值，右递归查找
            node.right = removeHelper(node.right, val);
        } else { // 待删除节点值=当前节点值，找到节点
            if (node.left == null || node.right == null) { // 如果当前节点的左节点或右节点为空
                TreeNode child = node.left != null ? node.left : node.right;
                if (child == null) { // 左右子树都为空时
                    return null;
                } else {
                    node = child;
                }
            } else { // 当前节点的左右子节点都不为空时，将中序遍历的下个节点删除，用该节点替换当前节点
                TreeNode tmp = node.right;
                while (tmp.left != null) {
                    tmp = tmp.left;
                }
                node.right = removeHelper(node.right, tmp.val);
                node.val = tmp.val;
            }
        }
        updateHeight(node); // 更新高度
        node = commonRotate(node); // 重新旋转，使得失衡树恢复平衡
        return node;
    }

    public static void main(String[] args) {
        testRightRotate();
        System.out.println("-------------------");
        testLeftRotate();
        System.out.println("-------------------");
        testCommonRotateLeft();
        System.out.println("-------------------");
        testCommonRotateRight();
        System.out.println("-------------------");
        AVLTree avlTree = new AVLTree();
        avlTree.testInsertAvlNode();
        System.out.println("-------------------");
        avlTree.testRemoveAvlNode(avlTree, 2);
        avlTree.testRemoveAvlNode(avlTree, 4);
    }

    // 验证删除节点后avl树是否平衡
    private void testRemoveAvlNode(AVLTree avlTree, int val) {
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        System.out.println("删除" + val +"前是否为AVL树？" + avlTreeValidatoar.isAVLTree(root));
        avlTree.remove(val);
        System.out.println("删除" + val +"后是否为AVL树？" + avlTreeValidatoar.isAVLTree(root));
    }

    // 验证添加节点后avl树是否平衡
    private void testInsertAvlNode() {
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        AVLTree avlTree = new AVLTree();
        avlTree.insert(2);
        avlTree.insert(1);
        avlTree.insert(4);
        avlTree.insert(3);
        avlTree.insert(5);
        avlTree.insert(6);
        System.out.println("是否为AVL树？" + avlTreeValidatoar.isAVLTree(avlTree.root));
//        return root;
    }

    // 验证右偏树
    private static void testCommonRotateRight() {
        TreeNode root = new TreeNode(2);
        TreeNode left = new TreeNode(1);
        TreeNode right = new TreeNode(3);
        root.left = left;
        root.right = right;
        TreeNode r_2_right = new TreeNode(5);
        right.right = r_2_right;
        TreeNode r_3_right = new TreeNode(4);
        r_2_right.left = r_3_right;
        // 更新每个节点的高度
        AVLTree avlTree = new AVLTree();
        avlTree.updateHeight(r_3_right);
        avlTree.updateHeight(r_2_right);
        avlTree.updateHeight(right);
        avlTree.updateHeight(left);
        avlTree.updateHeight(root);
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        System.out.println("是否为avl树？" + avlTreeValidatoar.isAVLTree(root));
        TreeNode balanceRoot = avlTree.commonRotate(root);
        System.out.println("先右旋后左旋是否为avl树？" + avlTreeValidatoar.isAVLTree(balanceRoot));
    }

    // 验证左偏树
    private static void testCommonRotateLeft() {
        TreeNode root = new TreeNode(4);
        TreeNode left = new TreeNode(3);
        TreeNode right = new TreeNode(5);
        root.left = left;
        root.right = right;
        TreeNode l_2_left = new TreeNode(1);
        left.left = l_2_left;
        TreeNode l_3_right = new TreeNode(2);
        l_2_left.right = l_3_right;
        AVLTree avlTree = new AVLTree();
        // 从底至上依次更新节点高度，如果不更新高度，则计算平衡因子会错误，导致执行结果不对
        avlTree.updateHeight(l_3_right);
        avlTree.updateHeight(l_2_left);
        avlTree.updateHeight(left);
        avlTree.updateHeight(right);
        avlTree.updateHeight(root);
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        System.out.println("是否为avl树？" + avlTreeValidatoar.isAVLTree(root));
        TreeNode balanceRoot = avlTree.commonRotate(root);
        System.out.println("先左旋后右旋是否为avl树？" + avlTreeValidatoar.isAVLTree(balanceRoot));
    }

    // 验证左旋
    private static void testLeftRotate() {
        TreeNode root = new TreeNode(1);
        TreeNode l_1_left = new TreeNode(0);
        TreeNode l_1_right = new TreeNode(4);
        root.left = l_1_left;
        root.right = l_1_right;
        TreeNode r_2_left = new TreeNode(3);
        TreeNode r_2_right = new TreeNode(5);
        l_1_right.left = r_2_left;
        l_1_right.right = r_2_right;
        TreeNode r_3_right = new TreeNode(6);
        r_2_right.right = r_3_right;
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        System.out.println("是否为avl树？" + avlTreeValidatoar.isAVLTree(root));
        AVLTree avlTree = new AVLTree();
        TreeNode balanceRoot = avlTree.leftRotate(root);
        System.out.println("左旋后是否为avl树？" + avlTreeValidatoar.isAVLTree(balanceRoot));
    }

    // 验证右旋
    private static void testRightRotate() {
        TreeNode root = new TreeNode(5);
        TreeNode l_1_left = new TreeNode(3);
        TreeNode l_1_right = new TreeNode(6);
        root.left = l_1_left;
        root.right = l_1_right;
        TreeNode l_2_left = new TreeNode(1);
        TreeNode l_2_right = new TreeNode(4);
        l_1_left.left = l_2_left;
        l_1_left.right = l_2_right;
        TreeNode l_3_left = new TreeNode(0);
        l_2_left.left = l_3_left;
        AVLTreeValidatoar avlTreeValidatoar = new AVLTreeValidatoar();
        System.out.println("是否为avl树？" + avlTreeValidatoar.isAVLTree(root));
        AVLTree avlTree = new AVLTree();
        TreeNode balanceRoot = avlTree.rightRotate(root);
        System.out.println("右旋后是否为avl树？" + avlTreeValidatoar.isAVLTree(balanceRoot));
    }

}
